#!/usr/bin/env python3
"""
读取 ROS2 点云和 IMU 数据，分批动画显示：
1. IMU加速度积分还原世界系轨迹（红色）
2. IMU角速度积分还原世界系姿态，并将点云转为世界系坐标（蓝色）
3. tqdm显示数据加载和动画进度，进度条描述显示IMU时间长度（HH:mm:ss.fff）
Usage:
    python3 plot_imu_cloud_world.py [--bag xxx.db3] [--bunch 20]
"""
import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions, StorageFilter
from sensor_msgs.msg import PointCloud2, Imu
from rclpy.serialization import deserialize_message
from tqdm import tqdm
from datetime import timedelta

# 四元数转旋转矩阵
def quat2R(qw, qx, qy, qz):
    R = np.array([
        [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
        [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)],
        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]
    ], dtype=np.float32)
    return R

# PointCloud2 转 xyz

def pc2_to_xyz(pc: PointCloud2):
    arr = np.frombuffer(pc.data, np.uint8).reshape(-1, pc.point_step)
    xyz = arr[:, :12].copy().view(np.float32).reshape(-1, 3)
    return xyz[np.isfinite(xyz).all(axis=1)]

# 加载并分批处理数据
def load_batches(bag_path, bunch_size):
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'), ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=['/unilidar/imu', '/unilidar/cloud']))

    # 统计总帧数和IMU时间范围
    total_msgs, imu_times = 0, []
    while reader.has_next():
        topic, data, t = reader.read_next()
        if topic == '/unilidar/imu':
            imu = deserialize_message(data, Imu)
            curr_time = imu.header.stamp.sec + imu.header.stamp.nanosec * 1e-9 if hasattr(imu, 'header') else t * 1e-9
            imu_times.append(curr_time)
        total_msgs += 1
    # IMU时间范围
    t_start = imu_times[0] if imu_times else 0
    t_end = imu_times[-1] if imu_times else 0
    # 重新打开reader
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'), ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=['/unilidar/imu', '/unilidar/cloud']))

    # IMU积分变量
    latest_R = np.eye(3, dtype=np.float32)
    latest_v = np.zeros(3, dtype=np.float32)
    latest_t = np.zeros(3, dtype=np.float32)
    last_time = None
    # 姿态积分变量（角速度）
    latest_q = np.array([1, 0, 0, 0], dtype=np.float32)  # w, x, y, z
    # 分批缓存
    batch_imu = []
    batch_cloud = []
    batches = []
    # tqdm进度条
    def fmt_time(secs):
        td = timedelta(seconds=secs)
        total_seconds = int(td.total_seconds())
        hours = total_seconds // 3600
        minutes = (total_seconds % 3600) // 60
        seconds = total_seconds % 60
        milliseconds = int(td.microseconds / 1000)
        return f'{hours:02d}:{minutes:02d}:{seconds:02d}.{milliseconds:03d}'
    pbar = tqdm(total=total_msgs, desc=f'加载数据 {fmt_time(t_end-t_start)}')

    # 记录每帧的IMU世界坐标和点云世界坐标
    imu_worlds, cloud_worlds = [], []
    while reader.has_next():
        topic, data, t = reader.read_next()
        if topic == '/unilidar/imu':
            imu = deserialize_message(data, Imu)
            # 时间戳
            curr_time = imu.header.stamp.sec + imu.header.stamp.nanosec * 1e-9 if hasattr(imu, 'header') else t * 1e-9
            if last_time is None:
                dt = 0.01
            else:
                dt = curr_time - last_time
            last_time = curr_time
            # 姿态（四元数）
            q = imu.orientation
            latest_R = quat2R(q.w, q.x, q.y, q.z)
            # 加速度（body系）
            acc_body = np.array([imu.linear_acceleration.x,
                                 imu.linear_acceleration.y,
                                 imu.linear_acceleration.z], dtype=np.float32)
            acc_world = latest_R @ acc_body
            # 积分速度和位置
            latest_v += acc_world * dt
            latest_t += latest_v * dt
            imu_worlds.append(latest_t.copy())
            batch_imu.append(latest_t.copy())
            # 角速度积分（简化：欧拉法，未归一化）
            wx, wy, wz = imu.angular_velocity.x, imu.angular_velocity.y, imu.angular_velocity.z
            omega = np.array([0, wx, wy, wz], dtype=np.float32)
            dq = 0.5 * quat_mult(latest_q, omega) * dt
            latest_q += dq
            latest_q /= np.linalg.norm(latest_q)
        elif topic == '/unilidar/cloud':
            pc = deserialize_message(data, PointCloud2)
            pts_imu = pc2_to_xyz(pc)
            # 用最新姿态和位置转到世界系
            R_cloud = quat2R(latest_q[0], latest_q[1], latest_q[2], latest_q[3])
            pts_world = (R_cloud @ pts_imu.T).T + latest_t
            cloud_worlds.append(pts_world)
            batch_cloud.append(pts_world)
        pbar.update(1)
        # 分批
        if len(batch_imu) >= bunch_size and len(batch_cloud) >= bunch_size:
            batches.append((np.vstack(batch_imu), batch_cloud[:bunch_size]))
            batch_imu, batch_cloud = [], batch_cloud[bunch_size:]
    # 尾部
    if batch_imu and batch_cloud:
        batches.append((np.vstack(batch_imu), batch_cloud))
    pbar.close()
    return batches, t_start, t_end

# 四元数乘法
def quat_mult(q, r):
    w0, x0, y0, z0 = q
    w1, x1, y1, z1 = r
    return np.array([
        w0*w1 - x0*x1 - y0*y1 - z0*z1,
        w0*x1 + x0*w1 + y0*z1 - z0*y1,
        w0*y1 - x0*z1 + y0*w1 + z0*x1,
        w0*z1 + x0*y1 - y0*x1 + z0*w1
    ], dtype=np.float32)

# 动画显示
def animate_batches(batches, t_start, t_end, interval=50):
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlabel('X world')
    ax.set_ylabel('Y world')
    ax.set_zlabel('Z world')
    ax.set_title('IMU & Cloud Trajectory (World Coordinates)')
    scat_imu = ax.scatter([], [], [], s=10, c='red', label='IMU traj')
    scat_cloud = ax.scatter([], [], [], s=1, c='blue', label='Cloud world')
    ax.legend(loc='upper right')
    # 全局范围
    all_imu = np.vstack([b[0] for b in batches]) if batches else np.empty((0, 3))
    all_cloud = np.vstack([np.vstack(b[1]) for b in batches]) if batches else np.empty((0, 3))
    if all_imu.shape[0] > 0 and all_cloud.shape[0] > 0:
        margin = 1.0
        mins = np.minimum(all_imu.min(axis=0), all_cloud.min(axis=0)) - margin
        maxs = np.maximum(all_imu.max(axis=0), all_cloud.max(axis=0)) + margin
        ax.set_xlim(mins[0], maxs[0])
        ax.set_ylim(mins[1], maxs[1])
        ax.set_zlim(mins[2], maxs[2])
    # tqdm动画进度
    def fmt_time(secs):
        td = timedelta(seconds=secs)
        total_seconds = int(td.total_seconds())
        hours = total_seconds // 3600
        minutes = (total_seconds % 3600) // 60
        seconds = total_seconds % 60
        milliseconds = int(td.microseconds / 1000)
        return f'{hours:02d}:{minutes:02d}:{seconds:02d}.{milliseconds:03d}'

    shorten_times = 4    
    filter_times = 5
    frames = int(len(batches)/filter_times/shorten_times)
    print(f"动画显示，共 {frames} 批次，IMU时间 {fmt_time(t_end-t_start)}")
    coords_imu = [np.vstack([batches[j][0] for j in range(i+1)]) for i in range(len(batches))]
    coords_cloud = [np.vstack([np.vstack(batches[j][1]) for j in range(i+1)]) for i in range(len(batches))]
    pbar = tqdm(total=frames, desc=f'动画 {fmt_time(t_end-t_start)}')

    def update(i):
        pts_imu = coords_imu[i*filter_times]
        pts_cloud = coords_cloud[i*filter_times]
        scat_imu._offsets3d = (pts_imu[:, 0], pts_imu[:, 1], pts_imu[:, 2])
        scat_cloud._offsets3d = (pts_cloud[:, 0], pts_cloud[:, 1], pts_cloud[:, 2])
        ax.set_title(f'{i+1}/{len(batches)} imu size:{pts_imu.shape[0]} cloud size:{pts_cloud.shape[0]}')
        pbar.n = i + 1
        pbar.set_description(f'动画 {fmt_time(t_end-t_start)}')
        pbar.refresh()
        if i + 1 == len(batches):
            pbar.close()
        return scat_imu, scat_cloud

    ani = animation.FuncAnimation(fig, update, frames=frames, interval=interval, blit=False, repeat=False)
    plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', default='/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3', help='ros2 bag db3 path')
    parser.add_argument('--bunch', type=int, default=20, help='frames per batch')
    parser.add_argument('--interval', type=int, default=50, help='ms per step')
    args = parser.parse_args()

    print('读取 bag 并分批加载 IMU/点云 世界系轨迹 …')
    batches, t_start, t_end = load_batches(args.bag, args.bunch)
    print(f'共 {len(batches)} 批')
    animate_batches(batches, t_start, t_end, interval=args.interval)
